import * as THREE from 'three';

class STLExporter {
    parse(scene, options = {}) {
        options = Object.assign({
            binary: false
        }, options);

        const binary = options.binary;

        //

        let output;
        let offset = 80; // STL二进制文件头
        let bufferLength;

        if (binary) {
            const triangles = 0;
            const lookup = new Map();
            const triangleCount = countTriangles(scene, lookup);
            bufferLength = triangleCount * 2 + triangleCount * 3 * 4 * 4 + 80 + 4;
            const arrayBuffer = new ArrayBuffer(bufferLength);
            output = new DataView(arrayBuffer);
            output.setUint32(offset, triangleCount, true); offset += 4;
        } else {
            output = '';
            output += 'solid exported\n';
        }

        scene.traverse(function (object) {
            if (object.isMesh) {
                const geometry = object.geometry;
                const matrixWorld = object.matrixWorld;

                if (geometry.isBufferGeometry) {
                    const vertices = geometry.getAttribute('position');
                    const indices = geometry.getIndex();

                    if (indices !== null) {
                        // 索引几何体
                        for (let i = 0; i < indices.count; i += 3) {
                            for (let j = 0; j < 3; j++) {
                                const vertexIndex = indices.getX(i + j);
                                processVertex(vertexIndex, vertices, matrixWorld);
                            }

                            if (!binary) output += '\tendloop\n\tendfacet\n';
                        }
                    } else {
                        // 非索引几何体
                        for (let i = 0; i < vertices.count; i += 3) {
                            for (let j = 0; j < 3; j++) {
                                processVertex(i + j, vertices, matrixWorld);
                            }

                            if (!binary) output += '\tendloop\n\tendfacet\n';
                        }
                    }
                }
            }
        });

        if (!binary) {
            output += 'endsolid exported\n';
        }

        return binary ? output.buffer : output;

        function processVertex(vertexIndex, vertexAttribute, matrixWorld) {
            const stride = vertexAttribute.itemSize;

            const x = vertexAttribute.getX(vertexIndex);
            const y = vertexAttribute.getY(vertexIndex);
            const z = vertexAttribute.getZ(vertexIndex);

            const vertex = new THREE.Vector3(x, y, z);
            vertex.applyMatrix4(matrixWorld);

            if (binary) {
                output.setFloat32(offset, vertex.x, true); offset += 4;
                output.setFloat32(offset, vertex.y, true); offset += 4;
                output.setFloat32(offset, vertex.z, true); offset += 4;
            } else {
                output += '\t\tvertex ' + vertex.x + ' ' + vertex.y + ' ' + vertex.z + '\n';
            }
        }

        function countTriangles(object, map) {
            let triangles = 0;

            object.traverse(function (obj) {
                if (!obj.isMesh) return;

                const geometry = obj.geometry;

                if (map.has(geometry)) {
                    triangles += map.get(geometry);
                } else {
                    let count = 0;

                    if (geometry.isBufferGeometry) {
                        const indices = geometry.getIndex();

                        if (indices !== null) {
                            count = indices.count / 3;
                        } else {
                            count = geometry.getAttribute('position').count / 3;
                        }
                    }

                    map.set(geometry, count);
                    triangles += count;
                }
            });

            return triangles;
        }
    }
}

export { STLExporter }; 